# -*- coding: utf-8 -*-
"""
@author: bijiahao
@date: 2025/10/20 17:43
@description:
"""
import json, asyncio
from fastapi import FastAPI
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
from openai import AsyncOpenAI
# from openai import OpenAI
from mcp.client.sse import sse_client
from mcp import ClientSession
import os

API_KEY = os.getenv("API_KEY")
app = FastAPI()
llm = AsyncOpenAI(api_key=API_KEY, base_url="https://dashscope.aliyuncs.com/compatible-mode/v1")
MODEL = "qwen-plus"
# llm = AsyncOpenAI(api_key=API_KEY, base_url="https://api.siliconflow.cn/v1")
# MODEL = "deepseek-ai/DeepSeek-V3"

# mcp服务地址
MCP_URL = "http://localhost:8001/sse"


class ChatReq(BaseModel):
    q: str


async def stream_answer(query: str):
    async with sse_client(MCP_URL) as (rd, wr), \
            ClientSession(rd, wr) as session:
        await session.initialize()
        tools = [{"type": "function",
                  "function": {"name": t.name, "description": t.description, "parameters": t.inputSchema}}
                 for t in (await session.list_tools()).tools]

        # 第一轮：让模型决定是否调用工具
        comp = await llm.chat.completions.create(
            model=MODEL, messages=[{"role": "user", "content": query}],
            tools=tools, tool_choice="auto", stream=False)
        msg = comp.choices[0].message

        if msg.tool_calls:  # 需要调用工具
            tool_call = msg.tool_calls[0]
            func = tool_call.function
            result = await session.call_tool(func.name, json.loads(func.arguments))
            # 把工具结果再扔给模型
            messages = [
                {"role": "user", "content": query},
                {"role": "assistant", "content": None, "tool_calls": [tool_call]},
                {"role": "tool", "content": result.content[0].text, "tool_call_id": tool_call.id}
            ]
        else:  # 无需工具
            yield msg.content or ""
            return

        # 第二轮：流式生成最终回答
        async for chunk in await llm.chat.completions.create(
                model=MODEL, messages=messages, stream=True):
            if tok := chunk.choices[0].delta.content:
                yield tok


@app.post("/chat")
def chat(req: ChatReq):
    return StreamingResponse(stream_answer(req.q), media_type="text/plain")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(
        "test_ai_chat:app",  # 应用导入路径
        host="0.0.0.0",  # 监听地址
        port=8002,  # 端口号
        reload=True,  # 开发时热重载
        log_level="info"  # 日志级别
    )
